# AprilTags Example
#
# This example shows the power of the OpenMV Cam to detect April Tags
# on the OpenMV Cam M7. The M4 versions cannot detect April Tags.

import sensor, image, time, math
from pid import PID
from pyb import Servo
from pyb import UART

#s1 = Servo(1) # P7
#s2 = Servo(2) # P8


sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QQVGA) # we run out of memory if the resolution is much bigger...
sensor.skip_frames(30)
sensor.set_auto_gain(False)  # must turn this off to prevent image washout...
sensor.set_auto_whitebal(False)  # must turn this off to prevent image washout...
clock = time.clock()

# 注意！与find_qrcodes不同，find_apriltags 不需要软件矫正畸变就可以工作。

# 注意，输出的姿态的单位是弧度，可以转换成角度，但是位置的单位是和你的大小有关，需要等比例换算

# f_x 是x的像素为单位的焦距。对于标准的OpenMV，应该等于2.8/3.984*656，这个值是用毫米为单位的焦距除以x方向的感光元件的长度，乘以x方向的感光元件的像素（OV7725）
# f_y 是y的像素为单位的焦距。对于标准的OpenMV，应该等于2.8/2.952*488，这个值是用毫米为单位的焦距除以y方向的感光元件的长度，乘以y方向的感光元件的像素（OV7725）

# c_x 是图像的x中心位置
# c_y 是图像的y中心位置

class CTRL_STATE:
    IDLE       = 0
    MOVE_NEXT  = 1
    SCAN_TAG   = 2
    FOLLOW     = 3
    FOCUS_TAG  = 4
    FIND_OVER  = 5
    FIND_FAIL  = 6


def degrees(radians):
    return (180 * radians) / math.pi

class Controller:
    def __init__(self):
        self.uart = UART(3, 19200)
        self.pan_servo=Servo(1)
        self.tilt_servo=Servo(2)
        self.pan_servo.calibration(500,2500,500)
        self.tilt_servo.calibration(500,2500,500)
        self.pan_servo.angle(90)
        self.tilt_servo.angle(90)


        self.steer_pos=[[0,90],[45,90],[90,90],[135,90],[180,90],\
                        [180,120],[135,120],[90,120],[45,120],[0,120]]
        self.steer_idx = -1
        self.scan_cnt  = 0
        self.state     = CTRL_STATE.IDLE
        self.lastTag   = None
        self.curImg    = None

        self.pan_pid   = PID(p=10,  i=1, imax=90 )
        self.tilt_pid  = PID(p=10,  i=1, imax=90 )
        self.invalidNum= 0

    def PidControl(self):
        dx = 2*self.lastTag.cx()/self.curImg.width()  - 1 # -1 1
        dy = 2*self.lastTag.cy()/self.curImg.height() - 1
        yaw = self.pan_servo.angle() + self.pan_pid.get_pid(dx,1)
        pitch = self.tilt_servo.angle() - self.tilt_pid.get_pid(dy,1)
        self.pan_servo.angle( yaw )
        self.tilt_servo.angle( pitch )
        return abs(dx)+abs(dy)

    def Move2NextPos(self):
        idx = (self.steer_idx+1)%len(self.steer_pos)
        self.Trans2PosByIndex(idx)
        self.steer_idx = idx

    def Trans2PosByIndex(self,idx):
        if idx != self.steer_idx:
            yaw = self.steer_pos[idx][0]
            pitch = self.steer_pos[idx][1]
            self.pan_servo.angle(yaw)
            self.tilt_servo.angle(pitch)
            time.sleep_ms(300)

    def GetImg(self,level):
        if level == 0:
            sensor.set_framesize(sensor.QQVGA)
        elif level == 1:
            sensor.set_framesize(sensor.QVGA)
        else:
            sensor.set_framesize(sensor.VGA)

        self.curImg = sensor.snapshot()

    def FindAprilTag(self,targetNum):
        img = self.curImg
        img_framesize_x=img.width()
        img_framesize_y=img.height()
        _fx = (2.8 / 3.984) * img_framesize_x # 默认值
        _fy = (2.8 / 2.952) * img_framesize_y # 默认值
        _cx = img_framesize_x * 0.5 # 默认值(image.w * 0.5)
        _cy = img_framesize_y * 0.5 # 默认值(image.h * 0.5)
        for tag in img.find_apriltags(fx=_fx, fy=_fy, cx=_cx, cy=_cy): # 默认为TAG36H11
            if tag.id() != targetNum:
                continue;
            img.draw_rectangle(tag.rect(), color = (255, 0, 0))
            img.draw_cross(tag.cx(), tag.cy(), color = (0, 255, 0))
            return tag
        return None


    def Run(self):
        if CTRL_STATE.IDLE == self.state:
            # TODO
            print("start"+ str(clock.fps()) )
            self.scan_cnt = 0
            self.state = CTRL_STATE.MOVE_NEXT

        elif CTRL_STATE.MOVE_NEXT == self.state:
            self.Move2NextPos()
            self.state = CTRL_STATE.SCAN_TAG

        elif CTRL_STATE.SCAN_TAG == self.state:
            self.GetImg(0)
            tag = self.FindAprilTag(3)
            if None == tag:
                self.scan_cnt = self.scan_cnt + 1
                if self.scan_cnt > 2*len(self.steer_pos):
                    self.state = CTRL_STATE.FIND_FAIL
                else:
                    self.state = CTRL_STATE.MOVE_NEXT
            else:
                self.lastTag    = tag
                self.invalidNum = 0
                self.state      = CTRL_STATE.FOLLOW

        elif CTRL_STATE.FOLLOW == self.state:
            err = self.PidControl()
            if err < 0.2:
                self.state = CTRL_STATE.FIND_OVER
            else:
                self.state = CTRL_STATE.FOCUS_TAG

        elif CTRL_STATE.FOCUS_TAG == self.state:
            if self.invalidNum > 5:
                self.state = CTRL_STATE.MOVE_NEXT
            else:
                self.GetImg(0)
                tag = self.FindAprilTag(3)
                if None == tag:
                    self.invalidNum = self.invalidNum + 1
                else:
                    self.invalidNum = 0
                    self.lastTag    = tag
                    self.state = CTRL_STATE.FOLLOW

        elif CTRL_STATE.FIND_OVER == self.state:
            ##TODO
            print("find")
            self.state = CTRL_STATE.IDLE
        elif CTRL_STATE.FIND_FAIL == self.state:
            ##TODO
            print("not find")
            self.state = CTRL_STATE.IDLE
ctrl = Controller()

while(True):
    clock.tick()
    ctrl.Run()

    #print(clock.fps())

